#pragma once
#include <unordered_map>
#include "glm.h"
#include <string>
#include "geometry.h"
#include "tgaimage.h"

class IShader;
class VertexShader;
class FragmentShader;

template <typename T>
struct UniformVar
{
    std::unordered_map<std::string, T> var;

    void add_var(const std::string& name)
    {
        var[name] = T();
    }
    T get_var(const std::string& name)
    {
        if (var.find(name) == var.end())
        {
            std::cout << "[glProgram] UniformVar get_var " << name << " not found" << std::endl;
            return T();
        }
        return var[name];
    }
    void set_var(const std::string& name, const T& value)
    {
        if (var.find(name) == var.end())
        {
            std::cout << "[glProgram] UniformVar set_var " << name << " not found" << std::endl;
            return;
        }
        var[name] = value;
    }
};

template <typename T>
struct VaryingVar
{
	std::unordered_map<std::string, std::array<T, 3>> var;

	void add_var(const std::string& name)
	{
		var[name] = std::array<T, 3>();
	}

	T get_var(const std::string& name, size_t index)
	{
		if (var.find(name) == var.end())
		{
			std::cout << "[glProgram] VaryingVar get_var " << name << " not found" << std::endl;
			return T();
		}
        if (index > 3)
        {
            std::cout << "[glProgram] VaryingVar get_var " << name << " index out of range" << std::endl;
            return T();
        }

		return var[name][index];
	}

    std::array<T, 3> get_var(const std::string& name)
	{
		if (var.find(name) == var.end())
		{
			std::cout << "[glProgram] VaryingVar get_var " << name << " not found" << std::endl;
			return T();
		}

		return var[name];
	}

	void set_var(const std::string& name,const T& value, size_t index)
	{
		if (var.find(name) == var.end())
		{
			std::cout << "[glProgram] VaryingVar set_var " << name << " not found" << std::endl;
			return;
		}
        if (index > 3)
        {
            std::cout << "[glProgram] VaryingVar set_var " << name << " index out of range" << std::endl;
            return;
        }
		var[name][index] = value;
	}
};

// 顶点所需的信息
struct VertexInfoBase
{
    size_t index;
    Vec3f location;

    virtual ~VertexInfoBase() = default;
};

class glProgram
{
    friend class VertexShader;
    friend class FragmentShader;
    friend class IShader;

public:
    glProgram();
    ~glProgram();

    void RegisterVertexShader(std::shared_ptr<VertexShader> vertexShader);// 注册顶点着色器
    void RegisterFragmentShader(std::shared_ptr<FragmentShader> fragmentShader);// 注册片段着色器
    void Draw(TGAImage& target, float* zbuffer, const std::array<std::shared_ptr<VertexInfoBase>, 3>& vertexInfos);// 绘制一个三角面
    void SetViewPort(int width, int height);// 设置视口
    glm::mat4 GetViewPort() const { return m_viewport; }

    template<typename T>
    T GetUniformVar(const std::string& name)
    {
        if constexpr (std::is_same_v<T, glm::mat4>)
        {
            return uniform_mat4.get_var(name);
        }
        else if constexpr (std::is_same_v<T, Vec4f>)
        {
            return uniform_vec4f.get_var(name);
        }
        else if constexpr (std::is_same_v<T, Vec3f>)
        {
            return uniform_vec3f.get_var(name);
        }
        else if constexpr (std::is_same_v<T,Vec2f>)
        {
            return uniform_vec2f.get_var(name);
        }
        else if constexpr (std::is_same_v<T, float>)
        {
            return uniform_float.get_var(name);
        }
        else if constexpr (std::is_same_v<T, int>)
        {
            return uniform_int.get_var(name);
        }
        else if constexpr (std::is_same_v<T, TGAImage*>)
        {
            return uniform_texture.get_var(name);
        }
    }

	template<typename T>
    void SetUniformVar(const std::string& name, const T& value)
    {
        if constexpr (std::is_same_v<T, glm::mat4>)
        {
            uniform_mat4.set_var(name, value);
        }
        else if constexpr (std::is_same_v<T, Vec4f>)
        {
            uniform_vec4f.set_var(name, value);
        }
        else if constexpr (std::is_same_v<T, Vec3f>)
        {
            uniform_vec3f.set_var(name, value);
        }
        else if constexpr (std::is_same_v<T, Vec2f>)
        {
            uniform_vec2f.set_var(name, value);
        }
        else if constexpr (std::is_same_v<T, float>)
        {
            uniform_float.set_var(name, value);
        }
        else if constexpr (std::is_same_v<T, int>)
        {
            uniform_int.set_var(name, value);
        }
        else if constexpr (std::is_same_v<T, TGAImage*>)
        {
            uniform_texture.set_var(name, value);
        }
    }

	template<typename T>
    void AddUniformVar(const std::string& name)
    {
        if constexpr (std::is_same_v<T, glm::mat4>)
        {
            uniform_mat4.add_var(name);
        }
        else if constexpr (std::is_same_v<T, Vec4f>)
        {
            uniform_vec4f.add_var(name);
        }
        else if constexpr (std::is_same_v<T, Vec3f>)
        {
            uniform_vec3f.add_var(name);
        }
        else if constexpr (std::is_same_v<T, Vec2f>)
        {
            uniform_vec2f.add_var(name);
        }
        else if constexpr (std::is_same_v<T, float>)
        {
            uniform_float.add_var(name);
        }
        else if constexpr (std::is_same_v<T, int>)
        {
            uniform_int.add_var(name);
        }
        else if constexpr (std::is_same_v<T, TGAImage*>)
        {
            uniform_texture.add_var(name);
        }
    }

    template<typename T>
    T GetVaryingVar(const std::string& name, size_t index)
    {
        if constexpr (std::is_same_v<T, Vec4f>)
        {
            return varying_vec4f.get_var(name, index);
        }
        else if constexpr (std::is_same_v<T, Vec3f>)
        {
            return varying_vec3f.get_var(name, index);
        }
        else if constexpr (std::is_same_v<T, Vec2f>)
        {
            return varying_vec2f.get_var(name, index);
        }
        else if constexpr (std::is_same_v<T, float>)
        {
            return varying_float.get_var(name, index);
        }
        else if constexpr (std::is_same_v<T, int>)
        {
            return varying_int.get_var(name, index);
        }
        else if constexpr (std::is_same_v<T, glm::mat4>)
        {
            return varying_mat4.get_var(name, index);
        }
		else if constexpr (std::is_same_v<T, glm::mat3>)
		{
			return varying_mat3.get_var(name, index);
		}
    }

	template<typename T>
	void SetVaryingVar(const std::string& name, const T& value, size_t index)
	{
        if constexpr (std::is_same_v<T, Vec4f>)
		{
            varying_vec4f.set_var(name, value, index);
		}
		else if constexpr (std::is_same_v<T, Vec3f>)
		{
            varying_vec3f.set_var(name, value, index);
		}
		else if constexpr (std::is_same_v<T, Vec2f>)
		{
            varying_vec2f.set_var(name, value, index);
		}
		else if constexpr (std::is_same_v<T, float>)
		{
            varying_float.set_var(name, value, index);
		}
		else if constexpr (std::is_same_v<T, int>)
		{
            varying_int.set_var(name, value, index);
		}
        else if constexpr (std::is_same_v<T, glm::mat4>)
		{
            varying_mat4.set_var(name, value, index);
		}
		else if constexpr (std::is_same_v<T, glm::mat3>)
		{
			varying_mat3.set_var(name, value, index);
		}
	}

	template<typename T>
	void AddVaryingVar(const std::string& name)
	{
        if constexpr (std::is_same_v<T, Vec4f>)
		{
            varying_vec4f.add_var(name);
		}
		else if constexpr (std::is_same_v<T, Vec3f>)
		{
            varying_vec3f.add_var(name);
		}
		else if constexpr (std::is_same_v<T, Vec2f>)
		{
            varying_vec2f.add_var(name);
		}
		else if constexpr (std::is_same_v<T, float>)
		{
            varying_float.add_var(name);
		}
		else if constexpr (std::is_same_v<T, int>)
		{
            varying_int.add_var(name);
		}
        else if constexpr (std::is_same_v<T, glm::mat4>)
		{
            varying_mat4.add_var(name);
		}
		else if constexpr (std::is_same_v<T, glm::mat3>)
		{
			varying_mat3.add_var(name);
		}
	}

protected:
    std::shared_ptr<VertexShader> m_vertexShader;
    std::shared_ptr<FragmentShader> m_fragmentShader;

    UniformVar<glm::mat4>   uniform_mat4;
	UniformVar<float>       uniform_float;
    UniformVar<Vec4f>       uniform_vec4f;
    UniformVar<Vec3f>       uniform_vec3f;
    UniformVar<Vec2f>       uniform_vec2f;
    UniformVar<int>         uniform_int;
    UniformVar<TGAImage*>   uniform_texture;

    VaryingVar<glm::mat4>   varying_mat4;
    VaryingVar<glm::mat3>   varying_mat3;
    VaryingVar<Vec4f>       varying_vec4f;
    VaryingVar<Vec3f>       varying_vec3f;
    VaryingVar<Vec2f>       varying_vec2f;
    VaryingVar<float>       varying_float;
    VaryingVar<int>         varying_int;

    glm::mat4 m_viewport;
};

